"""
数据记录模块

用于记录计算过程中的数据，包括：
1. 计算参数
2. 收敛历史
3. 物理量分布
4. 计算结果
"""

import os
import csv
import json
import numpy as np
from datetime import datetime
from .. import config as cfg

class DataLogger:
    """数据记录类"""
    
    def __init__(self, output_dir='output'):
        """
        初始化数据记录器
        
        参数：
        output_dir: str, 输出目录
        """
        self.output_dir = output_dir
        self.timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        # 创建本次计算的目录
        self.case_dir = os.path.join(output_dir, f'case_{self.timestamp}')
        os.makedirs(self.case_dir)
        
    def save_config(self):
        """保存配置参数"""
        config_file = os.path.join(self.case_dir, 'config.json')
        
        # 获取配置参数
        config_dict = {
            'grid': {
                'axial_nodes': cfg.AXIAL_NODES,
                'channel_length': cfg.CHANNEL_LENGTH
            },
            'geometry': {
                'channel_diameter': cfg.CHANNEL_DIAMETER
            },
            'inlet_conditions': {
                'pressure': cfg.INLET_PRESSURE,
                'temperature': cfg.INLET_TEMP,
                'mass_flux': cfg.MASS_FLUX
            },
            'channel_connections': cfg.CHANNEL_CONNECTIONS,
            'connection_coefficients': {str(k): v for k, v in cfg.CONNECTION_COEFFICIENTS.items()},
            'numerical_params': {
                'time_step': cfg.TIME_STEP,
                'max_iterations': cfg.MAX_ITERATIONS,
                'convergence_tolerance': cfg.CONVERGENCE_TOLERANCE
            },
            'physical_params': {
                'turbulent_mixing_coeff': cfg.TURBULENT_MIXING_COEFF,
                'bubble_drag_coeff': cfg.BUBBLE_DRAG_COEFF,
                'max_heat_flux': cfg.MAX_HEAT_FLUX
            },
            'material_params': {
                'cladding_density': cfg.CLADDING_DENSITY,
                'cladding_specific_heat': cfg.CLADDING_SPECIFIC_HEAT,
                'cladding_conductivity': cfg.CLADDING_CONDUCTIVITY
            }
        }
        
        # 保存为JSON文件
        with open(config_file, 'w') as f:
            json.dump(config_dict, f, indent=4)
            
    def save_convergence_history(self, history):
        """
        保存收敛历史
        
        参数：
        history: list, 包含每次迭代的误差信息
        """
        history_file = os.path.join(self.case_dir, 'convergence_history.csv')
        
        with open(history_file, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['Iteration', 'Max_Error', 'Mass_Imbalance', 'Energy_Imbalance'])
            for i, data in enumerate(history):
                writer.writerow([i+1, data['max_error'], data['mass_imbalance'], 
                               data['energy_imbalance']])
                
    def save_field_data(self, channels):
        """
        保存场数据
        
        参数：
        channels: dict, 子通道字典
        """
        # 创建场数据目录
        field_dir = os.path.join(self.case_dir, 'field_data')
        os.makedirs(field_dir)
        
        # 计算轴向位置
        z = np.linspace(0, cfg.CHANNEL_LENGTH, cfg.AXIAL_NODES)
        
        # 保存每个通道的数据
        for channel_id, channel in channels.items():
            channel_file = os.path.join(field_dir, f'channel_{channel_id}.csv')
            
            with open(channel_file, 'w', newline='') as f:
                writer = csv.writer(f)
                # 写入表头
                writer.writerow(['Position(m)', 'Alpha_l', 'Alpha_v', 'U_l(m/s)', 
                               'U_v(m/s)', 'V_l(m/s)', 'V_v(m/s)', 'T_l(K)', 
                               'T_v(K)', 'T_w(K)', 'P(Pa)'])
                               
                # 写入数据
                for i in range(cfg.AXIAL_NODES):
                    writer.writerow([z[i], channel.alpha_l[i], channel.alpha_v[i],
                                   channel.u_l[i], channel.u_v[i], channel.v_l[i],
                                   channel.v_v[i], channel.T_l[i], channel.T_v[i],
                                   channel.T_w[i], channel.pressure[i]])
                                   
    def save_results_summary(self, results):
        """
        保存计算结果摘要
        
        参数：
        results: dict, 计算结果字典
        """
        summary_file = os.path.join(self.case_dir, 'results_summary.json')
        
        # 转换numpy数组为列表
        for key, value in results.items():
            if isinstance(value, dict):
                for k, v in value.items():
                    if isinstance(v, np.ndarray):
                        results[key][k] = v.tolist()
            elif isinstance(value, np.ndarray):
                results[key] = value.tolist()
                
        # 保存为JSON文件
        with open(summary_file, 'w') as f:
            json.dump(results, f, indent=4)
            
    def save_performance_data(self, timing_data):
        """
        保存性能数据
        
        参数：
        timing_data: dict, 性能数据字典
        """
        performance_file = os.path.join(self.case_dir, 'performance.csv')
        
        with open(performance_file, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(['Module', 'Time(s)'])
            for module, time in timing_data.items():
                writer.writerow([module, time])
                
    def log_message(self, message):
        """
        记录日志消息
        
        参数：
        message: str, 日志消息
        """
        log_file = os.path.join(self.case_dir, 'calculation.log')
        
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        with open(log_file, 'a') as f:
            f.write(f'[{timestamp}] {message}\n') 